import os
from dotenv import load_dotenv 

import numpy as np
import pandas as pd

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_experimental.tools import PythonAstREPLTool
from langchain.chat_models import init_chat_model
from langchain_core.output_parsers import StrOutputParser
from langchain_core.output_parsers.openai_tools import JsonOutputKeyToolsParser

from langchain_core.runnables import RunnableLambda

from langchain_core.tools import tool

import requests
import json

from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser

class chatbot_with_tools:
    def __init__(self) -> None:
        load_dotenv(override=True)
        DeepSeek_API_KEY = os.getenv("DEEPSEEK_API_KEY")
        self.ds_api=DeepSeek_API_KEY
        #print(DeepSeek_API_KEY)  # 可以通过打印查看
        self.model = init_chat_model(model="deepseek-chat", model_provider="deepseek")  
        self.data_df=pd.read_csv(r'D:\Notes\LangChain\projects\data\blastchar\telco-customer-churn\telco-customer-churn\WA_Fn-UseC_-Telco-Customer-Churn.csv')
        #print(df.head())

    def simple_test(self, df):
        tool = PythonAstREPLTool(locals={"df": df})
        print(tool.invoke("df['SeniorCitizen'].mean()"))

    def simple_invoke(self, model, df):
        # 外部工具
        tool = PythonAstREPLTool(locals={"df": df})

        # 给model绑定外部工具
        llm_with_tools = model.bind_tools([tool])

        # 输入命令
        response = llm_with_tools.invoke(
            "我有一张表，名为'df'，请帮我计算MonthlyCharges字段的均值。"
        )
        
        print(response) # 这是一条调用外部工具的消息
    

    def invoke_with_parser(self, model, df):
        # 外部工具
        tool = PythonAstREPLTool(locals={"df": df})

        # 绑定外部工具
        llm_with_tools = model.bind_tools([tool])

        parser = JsonOutputKeyToolsParser(key_name=tool.name, first_tool_only=True)
        llm_chain = llm_with_tools | parser

        response=llm_chain.invoke("我有一张表，名为'df'，请帮我计算MonthlyCharges字段的均值。")
        print(response)
    
    def invoke_with_prompt(self, model, df):
        # 经过这一系列的工具链，这个方法包含官方内助工具的代码解释器
         # 外部工具
        tool = PythonAstREPLTool(locals={"df": df})

        # 绑定外部工具
        llm_with_tools = model.bind_tools([tool])


        # 提示词模板
        system = f"""
        你可以访问一个名为 `df` 的 pandas 数据框，你可以使用df.head().to_markdown() 查看数据集的基本信息， \
        请根据用户提出的问题，编写 Python 代码来回答。只返回代码，不返回其他内容。只允许使用 pandas 和内置库。
        """
        prompt = ChatPromptTemplate([
            ("system", system),
            ("user", "{question}")
        ])

        # 解析节点
        parser = JsonOutputKeyToolsParser(key_name=tool.name, first_tool_only=True)

        # 自定义操作方法
        def code_print(res):
            print("即将运行Python代码:", res['query'])
            return res
        
        print_node = RunnableLambda(code_print)

        code_chain = prompt | llm_with_tools | parser | print_node | tool 
        #response=code_chain.invoke("我有一张表，名为'df'，请帮我计算MonthlyCharges字段的均值。")
        #response2=code_chain.invoke({"question": "请帮我分析gender、SeniorCitizen和Churn三个字段之间的相关关系。"})
        response3=code_chain.invoke({"question": "请问MonthlyCharges取值最高的用户ID是？"})
        #print(response2)

    def invoke_with_self_define_tool(self, model, df):
        @tool
        def get_weather(loc):
            """
            查询即时天气函数
            :param loc: 必要参数，字符串类型，用于表示查询天气的具体城市名称，\
            注意，中国的城市需要用对应城市的英文名称代替，例如如果需要查询北京市天气，则loc参数需要输入'Beijing'；
            :return：OpenWeather API查询即时天气的结果，具体URL请求地址为：https://api.openweathermap.org/data/2.5/weather\
            返回结果对象类型为解析之后的JSON格式对象，并用字符串形式进行表示，其中包含了全部重要的天气信息
            """
            # Step 1.构建请求
            url = "https://api.openweathermap.org/data/2.5/weather"

            # Step 2.设置查询参数
            params = {
                "q": loc,               
                "appid": os.getenv("OPENWEATHER_API_KEY"),    # 输入API key
                "units": "metric",            # 使用摄氏度而不是华氏度
                "lang":"zh_cn"                # 输出语言为简体中文
            }

            # Step 3.发送GET请求
            response = requests.get(url, params=params)
            
            # Step 4.解析响应
            data = response.json()
            return json.dumps(data)

        tools = [get_weather]
        # 将工具绑定到模型
        llm_with_tools = model.bind_tools(tools)

        response = llm_with_tools.invoke("你好， 请问北京的天气怎么样？")
        print(response)

    def invoke_self_define_whole(self, model, df):
        # 这是一个完整的外部工具调用的方法
        @tool
        def get_weather(loc):
            """
            查询即时天气函数
            :param loc: 必要参数，字符串类型，用于表示查询天气的具体城市名称，\
            注意，中国的城市需要用对应城市的英文名称代替，例如如果需要查询北京市天气，则loc参数需要输入'Beijing'；
            :return：OpenWeather API查询即时天气的结果，具体URL请求地址为：https://api.openweathermap.org/data/2.5/weather\
            返回结果对象类型为解析之后的JSON格式对象，并用字符串形式进行表示，其中包含了全部重要的天气信息
            """
            # Step 1.构建请求
            url = "https://api.openweathermap.org/data/2.5/weather"

            # Step 2.设置查询参数
            params = {
                "q": loc,               
                "appid": os.getenv("OPENWEATHER_API_KEY"),    # 输入API key
                "units": "metric",            # 使用摄氏度而不是华氏度
                "lang":"zh_cn"                # 输出语言为简体中文
            }

            # Step 3.发送GET请求
            response = requests.get(url, params=params)
            
            # Step 4.解析响应
            data = response.json()
            return json.dumps(data)

        tools = [get_weather]

        # 将工具绑定到模型
        llm_with_tools = model.bind_tools(tools)

        parser = JsonOutputKeyToolsParser(key_name=get_weather.name, first_tool_only=True)

        # 拼接解析器、获取天气的方法
        get_weather_chain = llm_with_tools | parser | get_weather

        # Prompt 模板
        output_prompt = PromptTemplate.from_template(
            """你将收到一段 JSON 格式的天气数据，请用简洁自然的方式将其转述给用户。
        以下是天气 JSON 数据：

        ```json
        {weather_json}
        ```

        请将其转换为中文天气描述，例如：
        “北京当前天气晴，气温为 23°C，湿度 58%，风速 2.1 米/秒。”
        只返回一句话描述，不要其他说明或解释。"""
        )     

        output_chain = output_prompt | model | StrOutputParser()   

        # output_chain的简单测试
        '''
        weather_json = '{"coord": {"lon": 116.3972, "lat": 39.9075}, "weather": [{"id": 803, "main": "Clouds", "description": "\\u591a\\u4e91", "icon": "04d"}], "base": "stations", "main": {"temp": 34.94, "feels_like": 33.23, "temp_min": 34.94, "temp_max": 34.94, "pressure": 1002, "humidity": 22, "sea_level": 1002, "grnd_level": 997}, "visibility": 10000, "wind": {"speed": 6.23, "deg": 175, "gust": 9.33}, "clouds": {"all": 74}, "dt": 1749724384, "sys": {"type": 1, "id": 9609, "country": "CN", "sunrise": 1749674728, "sunset": 1749728589}, "timezone": 28800, "id": 1816670, "name": "Beijing", "cod": 200}'

        result = output_chain.invoke({"weather_json": weather_json})
        '''
        full_chain = get_weather_chain | output_chain
        response = full_chain.invoke("请问北京今天的天气如何？")
        print(response)


if __name__=='__main__':
    chatbot=chatbot_with_tools()
    #chatbot.invoke_with_prompt(chatbot.model, chatbot.data_df)
    
    # 一个简单的自定义方法调用
    #chatbot.invoke_with_self_define_tool(chatbot.model, chatbot.data_df)

    # 一个完整的自定义方法调用
    chatbot.invoke_self_define_whole(chatbot.model, chatbot.data_df)
